#!/usr/bin/env python
# -*- coding:utf-8 -*-

"""
=========================================
RabbitMQ 工具包
=========================================
"""

#Author : wu.zheng<wu.zheng@socialcredits.cn>


from functools import wraps
def get_queue(host, queue_name, user, password, port=5672, durable=True):
    """
     获得connection 和channel 连接对象

     params:
     -------
     host: MQ 主机IP
     queue_name: queue 名
     user:
     password:
     port:
     durable:  默认为True 消息会存磁盘, 否则不存
    """

    import pika
    credentials = pika.PlainCredentials(user , password)
    connection = pika.BlockingConnection(pika.ConnectionParameters(host, port, '/', credentials))
    channel = connection.channel()
    channel.queue_declare(queue=queue_name, durable=durable)
    return connection, channel

def get_queue_size(host, queue_name, user, password, port=15672):
    """
    获得queue 里消息数量

    params:
    -------

    host: MQ  主机IP
    queue_name : queue 名称,
    port : 端口 默认15672
    user:
    password
    """
    import requests
    from requests.auth import HTTPBasicAuth

    url = 'http://{host}:{port}/api/queues/%2f/{queue_name}'.format(host=host, port=port, queue_name=queue_name)
    response = requests.get(url, auth=HTTPBasicAuth(user, password))
    print response.content
    json_res = response.json()
    return json_res.get('messages', 0)



def rabbit_consumer(host, queue_name, user, password ,port=5672 ,durable=True):
    """
    RabbitMQ consumer 装饰器
    consumer 一定有一个参数叫message
    params :
    --------
    host:   MQ 主机IP
    queue_name: queue 名
    user:   user
    password: 密码
    port:    端口
    durable:  默认为True 消息会存磁盘,否则不存
    """
    connection, channel = get_queue(host, queue_name, user, password, port, durable)
    def decorator(function):
        def _decorator(*args, **kwargs):
            def callbak(ch, method, properties, body):
                item = eval(body)
                function(message=item, *args, **kwargs)
            channel.basic_qos(prefetch_count=1)
            channel.basic_consume(consumer_callback=callbak,
                      queue=queue_name,
                      no_ack=True)
            channel.start_consuming()
        return _decorator
    return decorator


def rabbit_producer(host, queue_name, user, password, port=5672, durable=True):
    """
    RabbitMq producer 装饰器

    被装饰的函数一定是一个可迭代的function

    param:
    --------
    host:   MQ 主机IP
    queue_name: queue 名
    user:   user
    password: 密码
    port:    端口
    durable:  默认为True 消息会存磁盘,否则不存
    """

    import pika
    connection, channel = get_queue(host, queue_name, user, password, port, durable)
    def decorator(function):
        @wraps(function)
        def _decorator(*args, **kwargs):
             for item in function(*args, **kwargs):
                 message = repr(item)
                 channel.basic_publish(
                     exchange='',
                     routing_key=queue_name,
                     body=message,
                     properties=pika.BasicProperties(delivery_mode=2)
                 )
             channel.close()
        return _decorator
    return decorator


